#!/usr/bin/env python3
import tkinter as tk
from tkinter import ttk
from tkinter import filedialog, messagebox
import os
import rosbag
from TopicsSelectGUI import TopicsSelect
import cv2
from cv_bridge import CvBridge
import json
from PIL import Image, ImageTk
from RemoteFileBrowser import RemoteFileBrowserDialog
import subprocess


class Bag2VideoGUI:
    def __init__(self):
        self.json_config_path = os.path.abspath(__file__)
        self.json_config_path = os.path.join(
            os.path.dirname(self.json_config_path), "config", "default_bag2video.json"
        )
        self.configuration = self.readconfig(self.json_config_path)
        self.bag_path = self.configuration["last_bag"]
        self.topics = self.configuration["selected_topics"]
        self.output_path = self.configuration["output_dir"]

        # Initialize GUI
        self.root = tk.Tk()
        self.root.title("ROS Bag to Video Converter")
        self.root.geometry("900x300")
        self.create_widgets()
        if not os.path.exists(self.bag_path):
            self.bag_path = ""
        if self.bag_path:
            self.bag = rosbag.Bag(self.bag_path)
            self.bag_path_entry.insert(0, self.bag_path)

    def create_widgets(self):
        # Menu bar
        self.menu_bar = tk.Menu(self.root)
        self.file_menu = tk.Menu(self.menu_bar, tearoff=0)
        self.file_menu.add_command(label="Open", command=self.open_file)
        self.file_menu.add_command(label="Remote", command=self.open_remote_browser)
        self.file_menu.add_command(label="Exit", command=self.root.quit)
        self.menu_bar.add_cascade(label="File", menu=self.file_menu)
        self.root.config(menu=self.menu_bar)

        # Bag path section
        self.bag_path_label = ttk.Label(self.root, text="Bag Path:")
        self.bag_path_label.grid(row=0, column=0, padx=10, pady=5, sticky="e")
        self.bag_path_entry = ttk.Entry(self.root, width=40)
        self.bag_path_entry.grid(row=0, column=1, padx=10, pady=5, sticky="w")
        # Progress bar
        self.process_bar = ttk.Progressbar(self.root, mode="determinate", length=300)
        self.process_bar.grid(row=1, column=0, columnspan=4, padx=10, pady=5)

        # Convert button
        self.convert_button = ttk.Button(
            self.root, text="Convert", command=self.convert
        )
        self.convert_button.grid(row=2, column=0, columnspan=3, pady=10)

        # Status label
        self.response_label = ttk.Label(self.root, text="", foreground="blue")
        self.response_label.grid(row=3, column=0, columnspan=3, pady=5)

        # Video Preview

        self.logo_path = os.path.abspath(__file__)
        self.logo_path = os.path.join(
            os.path.dirname(self.logo_path), "doc", "logo.png"
        )
        self.video_preview = ttk.Label(self.root)
        self.video_preview.grid(row=4, column=0, columnspan=3, pady=5)
        self.display_initial_image(self.logo_path)
        # Add padding only to grid-managed widgets
        for child in self.root.winfo_children():
            if isinstance(child, (ttk.Label, ttk.Entry, ttk.Button, ttk.Progressbar)):
                child.grid_configure(padx=5, pady=5)

    def display_initial_image(self, image_path):
        image = Image.open(image_path)
        img = ImageTk.PhotoImage(image)
        self.video_preview.imgtk = img
        self.video_preview.configure(image=img)

    def open_file(self):
        init_dir = os.path.dirname(self.bag_path) if self.bag_path else os.getcwd()
        file = filedialog.askopenfilename(
            filetypes=[("ROS Bag files", "*.bag")],
            initialdir=init_dir,
            initialfile=self.bag_path,
        )
        if not file:
            return
        self.bag_path_entry.delete(0, tk.END)
        self.bag_path_entry.insert(0, file)
        self.bag_path = file
        self.bag = rosbag.Bag(self.bag_path)
        topics_info = self.bag.get_type_and_topic_info()
        topics = [
            topic
            for topic, info in topics_info.topics.items()
            if info.msg_type == "sensor_msgs/Image"
        ]
        topics_window = TopicsSelect(
            self.root, topics, self.topics, self.topic_select_callback
        )
        topics_window.grab_set()

    def topic_select_callback(self, selected_topics):
        self.topics = selected_topics

    def open_remote_browser(self):
        remote_browser = RemoteFileBrowserDialog(
            self.root, "bag", "/tmp", res_callback=self.remote_browser_callback
        )
        remote_browser.grab_set()

    def remote_browser_callback(self, result):
        if result:
            self.bag_path_entry.delete(0, tk.END)
            self.bag_path_entry.insert(0, result)
            self.bag_path = result
            self.bag = rosbag.Bag(self.bag_path)
            topics_info = self.bag.get_type_and_topic_info()
            topics = [
                topic
                for topic, info in topics_info.topics.items()
                if info.msg_type == "sensor_msgs/Image"
            ]
            topics_window = TopicsSelect(
                self.root, topics, self.topics, self.topic_select_callback
            )
            topics_window.grab_set()

    def convert(self):
        if not self.topics:
            messagebox.showerror("Error", "Please select topics.")
            return
        init_path = os.path.dirname(self.bag_path) if self.bag_path else os.getcwd()
        output_path = filedialog.askdirectory(
            initialdir=init_path, title="Select output directory"
        )
        if not output_path:
            return
        base_name = os.path.splitext(os.path.basename(self.bag_path))[0]
        output_path = os.path.join(output_path, base_name)
        os.makedirs(output_path, exist_ok=True)

        for i, topic in enumerate(self.topics):
            video_file_name = topic.replace("/", "_").lstrip("_") + ".mp4"
            video_file = os.path.join(output_path, video_file_name)
            self.process_bar["value"] = 0
            self.response_label["text"] = (
                f"Converting {topic} ({i+1}/{len(self.topics)})..."
            )
            self.root.update()
            self.rosbag_to_video(self.bag, video_file, topic)

        self.response_label["text"] = "Convert finished."
        self.saveconfig(self.json_config_path)
        self.display_initial_image(self.logo_path)
        messagebox.showinfo("Info", f"Convert finished. Videos saved in {output_path}")

    def rosbag_to_video(
        self,
        bag,
        video_file,
        image_topic,
    ):
        bridge = CvBridge()

        fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # 使用 MP4 编码

        bag_start_time = bag.get_start_time()
        bag_end_time = bag.get_end_time()
        frames = 0
        topic_start_time = -1
        topic_end_time = 0
        width = 0
        height = 0
        for topic, msg, t in bag.read_messages(topics=[image_topic]):
            time = msg.header.stamp.to_sec()
            if topic_start_time == -1:
                topic_start_time = time
            topic_end_time = time
            frames += 1
            if width == 0 and height == 0:
                width = msg.width
                height = msg.height
        fps = frames / (topic_end_time - topic_start_time)
        print(f"FPS: {fps}")
        video_file_tmp = video_file.replace(".mp4", "_tmp.mp4")
        video_writer = cv2.VideoWriter(video_file_tmp, fourcc, fps, (width, height))
        for topic, msg, t in bag.read_messages(topics=[image_topic]):
            progress = int(
                (t.to_sec() - bag_start_time) / (bag_end_time - bag_start_time) * 100
            )
            self.process_bar["value"] = progress
            self.root.update()
            try:
                if msg.encoding == "bgr8":
                    cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
                    image_resized = cv2.resize(cv_image, (width, height))
                    video_writer.write(image_resized)
                elif msg.encoding == "rgb8":
                    cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
                    cv_image = cv2.cvtColor(cv_image, cv2.COLOR_RGB2BGR)
                    image_resized = cv2.resize(cv_image, (width, height))
                    video_writer.write(image_resized)
                elif msg.encoding == "mono8":
                    cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
                    image_resized = cv2.resize(cv_image, (width, height))
                    video_writer.write(image_resized)
                elif msg.encoding == "16UC1" or msg.encoding == "32FC1" or msg.encoding == "mono16":
                    cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
                    if msg.encoding == "16UC1" or msg.encoding == "mono16":
                        cv_image = cv2.convertScaleAbs(cv_image, alpha=0.03)
                    elif msg.encoding == "32FC1" :
                        cv_image = cv2.convertScaleAbs(cv_image, alpha=255.0 / 5)
                    cv_image = cv2.applyColorMap(cv_image, cv2.COLORMAP_JET)
                    depth_colormap_resized = cv2.resize(cv_image, (width, height))
                    video_writer.write(depth_colormap_resized)
                cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
                cv_image = cv2.resize(cv_image, (int(width / 10), int(height / 10)))
                img = ImageTk.PhotoImage(Image.fromarray(cv_image))
                self.video_preview.configure(image=img)
            except Exception as e:
                print(f"Failed to process image: {e}")

        video_writer.release()
        self.mp4_to_H264(video_file_tmp, video_file, fps=fps)

    def mp4_to_H264(self, mp4_file, output_file, fps):
        command = f"ffmpeg -i {mp4_file} -c:v libx264 -crf 23 -c:a aac -strict experimental -r {fps} -b:a 192k -ac 2 {output_file}"
        subprocess.run(command, shell=True)
        os.remove(mp4_file)

    def readconfig(self, config_path):
        try:
            with open(config_path, "r") as file:
                return json.load(file)
        except FileNotFoundError:
            self.create_config_json(config_path)
            with open(config_path, "r") as file:
                return json.load(file)

    def create_config_json(self, config_path):
        os.makedirs(os.path.dirname(config_path), exist_ok=True)
        default_config = {"last_bag": "", "selected_topics": [], "output_dir": ""}
        with open(config_path, "w") as file:
            json.dump(default_config, file, indent=4)

    def saveconfig(self, config_path):
        self.configuration.update(
            {
                "last_bag": self.bag_path,
                "selected_topics": self.topics,
                "output_dir": self.output_path,
            }
        )
        with open(config_path, "w") as file:
            json.dump(self.configuration, file, indent=4)


if __name__ == "__main__":
    app = Bag2VideoGUI()
    app.root.mainloop()
